#!/usr/bin/env python3.8
# -*- coding: utf-8 -*-
"""
YOLO数据集标注工具
使用tkinter制作的GUI界面，支持图像标注和YOLO格式输出
"""

import tkinter as tk
from tkinter import ttk, filedialog, messagebox, simpledialog
from PIL import Image, ImageTk
import os
import json
import math
from typing import List, Tuple, Dict, Optional


class YOLOAnnotator:
    def __init__(self, root):
        self.root = root
        self.root.title("密涅瓦yolo数据标注工具")
        self.root.geometry("1200x800")
        
        # 数据存储
        self.current_image_path = None
        self.current_image = None
        self.photo = None
        self.annotations = []  # 存储当前图像的标注
        self.classes = []  # 空类别列表，用户需要手动添加
        self.current_class = 0
        self.drawing = False
        self.start_x = 0
        self.start_y = 0
        self.current_bbox = None
        
        # 多段线标注相关
        self.annotation_mode = "rectangle"  # "rectangle" 或 "polygon"
        self.polygon_points = []  # 存储多边形点
        self.current_polygon = None  # 当前正在绘制的多边形
        
        # 创建界面
        self.create_widgets()
        self.setup_bindings()
        
    def create_widgets(self):
        """创建GUI组件"""
        # 主框架
        main_frame = ttk.Frame(self.root)
        main_frame.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)
        
        # 左侧控制面板
        control_frame = ttk.Frame(main_frame)
        control_frame.pack(side=tk.LEFT, fill=tk.Y, padx=(0, 5))
        
        # 文件操作
        file_frame = ttk.LabelFrame(control_frame, text="文件操作")
        file_frame.pack(fill=tk.X, pady=(0, 5))
        
        ttk.Button(file_frame, text="打开图像", command=self.open_image).pack(fill=tk.X, pady=2)
        ttk.Button(file_frame, text="打开文件夹", command=self.open_folder).pack(fill=tk.X, pady=2)
        ttk.Button(file_frame, text="保存标注", command=self.save_annotations).pack(fill=tk.X, pady=2)
        ttk.Button(file_frame, text="加载标注", command=self.load_annotations).pack(fill=tk.X, pady=2)
        
        # 类别管理
        class_frame = ttk.LabelFrame(control_frame, text="类别管理")
        class_frame.pack(fill=tk.X, pady=(0, 5))
        
        # 类别选择
        ttk.Label(class_frame, text="当前类别:").pack(anchor=tk.W)
        self.class_var = tk.StringVar(value="请先添加类别")
        self.class_combo = ttk.Combobox(class_frame, textvariable=self.class_var, 
                                       values=self.classes, state="readonly")
        self.class_combo.pack(fill=tk.X, pady=2)
        self.class_combo.bind("<<ComboboxSelected>>", self.on_class_change)
        
        # 添加类别
        add_class_frame = ttk.Frame(class_frame)
        add_class_frame.pack(fill=tk.X, pady=2)
        self.new_class_var = tk.StringVar()
        ttk.Entry(add_class_frame, textvariable=self.new_class_var).pack(side=tk.LEFT, fill=tk.X, expand=True)
        ttk.Button(add_class_frame, text="添加", command=self.add_class).pack(side=tk.RIGHT, padx=(2, 0))
        
        # 标注模式选择
        mode_frame = ttk.LabelFrame(control_frame, text="标注模式")
        mode_frame.pack(fill=tk.X, pady=(0, 5))
        
        self.mode_var = tk.StringVar(value="rectangle")
        ttk.Radiobutton(mode_frame, text="矩形框", variable=self.mode_var, 
                       value="rectangle", command=self.on_mode_change).pack(anchor=tk.W)
        ttk.Radiobutton(mode_frame, text="多边形", variable=self.mode_var, 
                       value="polygon", command=self.on_mode_change).pack(anchor=tk.W)
        
        # 标注操作
        annotate_frame = ttk.LabelFrame(control_frame, text="标注操作")
        annotate_frame.pack(fill=tk.X, pady=(0, 5))
        
        ttk.Button(annotate_frame, text="删除选中", command=self.delete_selected).pack(fill=tk.X, pady=2)
        ttk.Button(annotate_frame, text="清空所有", command=self.clear_annotations).pack(fill=tk.X, pady=2)
        ttk.Button(annotate_frame, text="完成多边形", command=self.finish_polygon).pack(fill=tk.X, pady=2)
        
        # 标注列表
        list_frame = ttk.LabelFrame(control_frame, text="标注列表")
        list_frame.pack(fill=tk.BOTH, expand=True)
        
        # 创建Treeview显示标注
        columns = ("ID", "类别", "坐标")
        self.annotation_tree = ttk.Treeview(list_frame, columns=columns, show="headings", height=10)
        
        for col in columns:
            self.annotation_tree.heading(col, text=col)
            self.annotation_tree.column(col, width=80)
        
        # 滚动条
        scrollbar = ttk.Scrollbar(list_frame, orient=tk.VERTICAL, command=self.annotation_tree.yview)
        self.annotation_tree.configure(yscrollcommand=scrollbar.set)
        
        self.annotation_tree.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
        scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
        
        # 绑定选择事件
        self.annotation_tree.bind("<<TreeviewSelect>>", self.on_annotation_select)
        
        # 右侧图像显示区域
        image_frame = ttk.Frame(main_frame)
        image_frame.pack(side=tk.RIGHT, fill=tk.BOTH, expand=True)
        
        # 图像画布
        self.canvas = tk.Canvas(image_frame, bg="white", cursor="crosshair")
        self.canvas.pack(fill=tk.BOTH, expand=True)
        
        # 状态栏
        self.status_var = tk.StringVar(value="就绪")
        status_bar = ttk.Label(self.root, textvariable=self.status_var, relief=tk.SUNKEN)
        status_bar.pack(side=tk.BOTTOM, fill=tk.X)
        
    def setup_bindings(self):
        """设置键盘和鼠标绑定"""
        # 鼠标事件
        self.canvas.bind("<Button-1>", self.on_mouse_press)
        self.canvas.bind("<B1-Motion>", self.on_mouse_drag)
        self.canvas.bind("<ButtonRelease-1>", self.on_mouse_release)
        self.canvas.bind("<Button-3>", self.on_right_click)  # 右键完成多边形
        
        # 键盘事件
        self.root.bind("<Key>", self.on_key_press)
        self.root.focus_set()
        
    def open_image(self):
        """打开图像文件"""
        file_path = filedialog.askopenfilename(
            title="选择图像文件",
            filetypes=[("图像文件", "*.jpg *.jpeg *.png *.bmp *.tiff"), ("所有文件", "*.*")]
        )
        if file_path:
            self.load_image(file_path)
            
    def open_folder(self):
        """打开图像文件夹"""
        folder_path = filedialog.askdirectory(title="选择图像文件夹")
        if folder_path:
            self.image_folder = folder_path
            self.image_files = [f for f in os.listdir(folder_path) 
                              if f.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.tiff'))]
            if self.image_files:
                self.current_image_index = 0
                self.load_image(os.path.join(folder_path, self.image_files[0]))
                self.status_var.set(f"已加载文件夹，共{len(self.image_files)}张图像")
            else:
                messagebox.showwarning("警告", "文件夹中没有找到图像文件")
                
    def load_image(self, image_path):
        """加载图像到画布"""
        try:
            self.current_image_path = image_path
            self.current_image = Image.open(image_path)
            
            # 计算缩放比例以适应画布
            canvas_width = self.canvas.winfo_width()
            canvas_height = self.canvas.winfo_height()
            
            if canvas_width <= 1 or canvas_height <= 1:
                # 如果画布还没有初始化，使用默认大小
                canvas_width, canvas_height = 800, 600
                
            img_width, img_height = self.current_image.size
            
            # 计算缩放比例，保持宽高比
            scale_x = canvas_width / img_width
            scale_y = canvas_height / img_height
            self.scale = min(scale_x, scale_y, 1.0)  # 不放大图像
            
            # 计算显示尺寸
            self.display_width = int(img_width * self.scale)
            self.display_height = int(img_height * self.scale)
            
            # 缩放图像
            self.display_image = self.current_image.resize((self.display_width, self.display_height), Image.LANCZOS)
            self.photo = ImageTk.PhotoImage(self.display_image)
            
            # 清除画布并显示图像
            self.canvas.delete("all")
            self.canvas.create_image(0, 0, anchor=tk.NW, image=self.photo)
            
            # 加载对应的标注文件
            self.load_annotations_for_image()
            
            self.status_var.set(f"已加载: {os.path.basename(image_path)}")
            
        except Exception as e:
            messagebox.showerror("错误", f"无法加载图像: {str(e)}")
            
    def load_annotations_for_image(self):
        """加载当前图像的标注"""
        if not self.current_image_path:
            return
            
        # 查找对应的标注文件
        base_name = os.path.splitext(self.current_image_path)[0]
        annotation_file = base_name + ".txt"
        
        self.annotations = []
        if os.path.exists(annotation_file):
            try:
                with open(annotation_file, 'r', encoding='utf-8') as f:
                    for line in f:
                        parts = line.strip().split()
                        if len(parts) >= 5:
                            class_id = int(parts[0])
                            
                            # 判断是矩形还是多边形
                            if len(parts) == 5:  # 矩形格式：class_id x_center y_center width height
                                x_center = float(parts[1])
                                y_center = float(parts[2])
                                width = float(parts[3])
                                height = float(parts[4])
                                
                                # 转换为像素坐标
                                img_width, img_height = self.current_image.size
                                x1 = (x_center - width/2) * img_width
                                y1 = (y_center - height/2) * img_height
                                x2 = (x_center + width/2) * img_width
                                y2 = (y_center + height/2) * img_height
                                
                                self.annotations.append({
                                    'class_id': class_id,
                                    'class_name': self.classes[class_id] if class_id < len(self.classes) else f"class_{class_id}",
                                    'bbox': (x1, y1, x2, y2),
                                    'type': 'rectangle'
                                })
                                
                            else:  # 多边形格式：class_id x1 y1 x2 y2 ... xn yn
                                points = []
                                for i in range(1, len(parts), 2):
                                    if i + 1 < len(parts):
                                        x = float(parts[i]) * img_width
                                        y = float(parts[i + 1]) * img_height
                                        points.append((x, y))
                                
                                if len(points) >= 3:
                                    self.annotations.append({
                                        'class_id': class_id,
                                        'class_name': self.classes[class_id] if class_id < len(self.classes) else f"class_{class_id}",
                                        'points': points,
                                        'type': 'polygon'
                                    })
            except Exception as e:
                messagebox.showerror("错误", f"无法加载标注文件: {str(e)}")
                
        self.update_annotation_list()
        self.draw_annotations()
        
    def on_mode_change(self):
        """标注模式切换"""
        self.annotation_mode = self.mode_var.get()
        if self.annotation_mode == "polygon":
            self.canvas.configure(cursor="crosshair")
            self.status_var.set("多边形模式：左键添加点，右键完成")
        else:
            self.canvas.configure(cursor="crosshair")
            self.status_var.set("矩形框模式：拖拽绘制矩形")
            
    def on_mouse_press(self, event):
        """鼠标按下事件"""
        if not self.current_image:
            return
            
        if self.annotation_mode == "rectangle":
            self.drawing = True
            self.start_x = event.x
            self.start_y = event.y
            # 删除当前正在绘制的矩形
            self.canvas.delete("current_bbox")
        elif self.annotation_mode == "polygon":
            # 添加多边形点
            self.polygon_points.append((event.x, event.y))
            self.draw_polygon_preview()
        
    def on_mouse_drag(self, event):
        """鼠标拖拽事件"""
        if not self.drawing or not self.current_image or self.annotation_mode != "rectangle":
            return
            
        # 删除之前的临时矩形
        self.canvas.delete("current_bbox")
        
        # 绘制新的临时矩形
        self.canvas.create_rectangle(
            self.start_x, self.start_y, event.x, event.y,
            outline="red", width=2, tags="current_bbox"
        )
        
    def on_mouse_release(self, event):
        """鼠标释放事件"""
        if not self.drawing or not self.current_image or self.annotation_mode != "rectangle":
            return
            
        self.drawing = False
        
        # 删除临时矩形
        self.canvas.delete("current_bbox")
        
        # 计算矩形坐标
        x1, y1 = min(self.start_x, event.x), min(self.start_y, event.y)
        x2, y2 = max(self.start_x, event.x), max(self.start_y, event.y)
        
        # 检查矩形大小
        if abs(x2 - x1) < 5 or abs(y2 - y1) < 5:
            return
            
        # 转换为图像坐标
        img_x1 = x1 / self.scale
        img_y1 = y1 / self.scale
        img_x2 = x2 / self.scale
        img_y2 = y2 / self.scale
        
        # 添加标注
        annotation = {
            'class_id': self.current_class,
            'class_name': self.classes[self.current_class],
            'bbox': (img_x1, img_y1, img_x2, img_y2),
            'type': 'rectangle'
        }
        self.annotations.append(annotation)
        
        # 更新界面
        self.update_annotation_list()
        self.draw_annotations()
        
    def on_right_click(self, event):
        """右键完成多边形"""
        if self.annotation_mode == "polygon" and len(self.polygon_points) >= 3:
            self.finish_polygon()
            
    def draw_polygon_preview(self):
        """绘制多边形预览"""
        if len(self.polygon_points) < 2:
            return
            
        # 删除之前的多边形预览
        self.canvas.delete("polygon_preview")
        
        # 绘制多边形预览
        points = []
        for x, y in self.polygon_points:
            points.extend([x, y])
            
        if len(points) >= 6:  # 至少3个点
            self.canvas.create_polygon(points, outline="red", width=2, 
                                     fill="", tags="polygon_preview")
            
        # 绘制点
        for x, y in self.polygon_points:
            self.canvas.create_oval(x-3, y-3, x+3, y+3, fill="red", 
                                  outline="red", tags="polygon_preview")
            
    def finish_polygon(self):
        """完成多边形标注"""
        if len(self.polygon_points) < 3:
            messagebox.showwarning("警告", "多边形至少需要3个点")
            return
            
        # 转换为图像坐标
        img_points = []
        for x, y in self.polygon_points:
            img_x = x / self.scale
            img_y = y / self.scale
            img_points.append((img_x, img_y))
            
        # 添加多边形标注
        annotation = {
            'class_id': self.current_class,
            'class_name': self.classes[self.current_class],
            'points': img_points,
            'type': 'polygon'
        }
        self.annotations.append(annotation)
        
        # 清理多边形点
        self.polygon_points = []
        self.canvas.delete("polygon_preview")
        
        # 更新界面
        self.update_annotation_list()
        self.draw_annotations()
        self.status_var.set("多边形标注完成")
        
    def draw_annotations(self):
        """绘制所有标注"""
        if not self.current_image:
            return
            
        # 清除之前的标注
        self.canvas.delete("annotation")
        
        for i, annotation in enumerate(self.annotations):
            class_name = annotation['class_name']
            annotation_type = annotation.get('type', 'rectangle')
            
            if annotation_type == 'rectangle':
                bbox = annotation['bbox']
                # 转换为显示坐标
                x1 = bbox[0] * self.scale
                y1 = bbox[1] * self.scale
                x2 = bbox[2] * self.scale
                y2 = bbox[3] * self.scale
                
                # 绘制矩形框
                rect_id = self.canvas.create_rectangle(
                    x1, y1, x2, y2,
                    outline="red", width=2, tags=("annotation", f"annotation_{i}")
                )
                
                # 绘制类别标签
                label_id = self.canvas.create_text(
                    x1, y1-10, anchor=tk.NW,
                    text=class_name, fill="red", font=("Arial", 10, "bold"),
                    tags=("annotation", f"annotation_{i}")
                )
                
            elif annotation_type == 'polygon':
                points = annotation['points']
                # 转换为显示坐标
                display_points = []
                for x, y in points:
                    display_x = x * self.scale
                    display_y = y * self.scale
                    display_points.extend([display_x, display_y])
                
                if len(display_points) >= 6:  # 至少3个点
                    # 绘制多边形
                    poly_id = self.canvas.create_polygon(
                        display_points, outline="blue", width=2, fill="",
                        tags=("annotation", f"annotation_{i}")
                    )
                    
                    # 绘制类别标签（在多边形中心）
                    if len(points) > 0:
                        center_x = sum(p[0] for p in points) / len(points) * self.scale
                        center_y = sum(p[1] for p in points) / len(points) * self.scale
                        label_id = self.canvas.create_text(
                            center_x, center_y, anchor=tk.CENTER,
                            text=class_name, fill="blue", font=("Arial", 10, "bold"),
                            tags=("annotation", f"annotation_{i}")
                        )
            
    def update_annotation_list(self):
        """更新标注列表"""
        # 清除现有项目
        for item in self.annotation_tree.get_children():
            self.annotation_tree.delete(item)
            
        # 添加标注项目
        for i, annotation in enumerate(self.annotations):
            annotation_type = annotation.get('type', 'rectangle')
            
            if annotation_type == 'rectangle':
                bbox = annotation['bbox']
                coord_str = f"({int(bbox[0])},{int(bbox[1])},{int(bbox[2])},{int(bbox[3])})"
            elif annotation_type == 'polygon':
                points = annotation['points']
                coord_str = f"多边形({len(points)}个点)"
            else:
                coord_str = "未知类型"
                
            self.annotation_tree.insert("", "end", values=(
                i,
                annotation['class_name'],
                coord_str
            ))
            
    def on_class_change(self, event):
        """类别选择改变"""
        selected_class = self.class_var.get()
        if selected_class in self.classes:
            self.current_class = self.classes.index(selected_class)
            
    def add_class(self):
        """添加新类别"""
        new_class = self.new_class_var.get().strip()
        if new_class and new_class not in self.classes:
            self.classes.append(new_class)
            self.class_combo['values'] = self.classes
            self.class_var.set(new_class)
            self.current_class = len(self.classes) - 1
            self.new_class_var.set("")
            self.status_var.set(f"已添加类别: {new_class}")
        elif new_class in self.classes:
            messagebox.showwarning("警告", "类别已存在")
        else:
            messagebox.showwarning("警告", "请输入有效的类别名称")
            
    def on_annotation_select(self, event):
        """标注选择事件"""
        selection = self.annotation_tree.selection()
        if selection:
            item = self.annotation_tree.item(selection[0])
            annotation_id = int(item['values'][0])
            self.highlight_annotation(annotation_id)
            
    def highlight_annotation(self, annotation_id):
        """高亮显示选中的标注"""
        # 清除之前的高亮
        self.canvas.delete("highlight")
        
        if 0 <= annotation_id < len(self.annotations):
            annotation = self.annotations[annotation_id]
            bbox = annotation['bbox']
            
            # 转换为显示坐标
            x1 = bbox[0] * self.scale
            y1 = bbox[1] * self.scale
            x2 = bbox[2] * self.scale
            y2 = bbox[3] * self.scale
            
            # 绘制高亮矩形
            self.canvas.create_rectangle(
                x1, y1, x2, y2,
                outline="yellow", width=3, tags="highlight"
            )
            
    def delete_selected(self):
        """删除选中的标注"""
        selection = self.annotation_tree.selection()
        if selection:
            item = self.annotation_tree.item(selection[0])
            annotation_id = int(item['values'][0])
            
            if 0 <= annotation_id < len(self.annotations):
                del self.annotations[annotation_id]
                self.update_annotation_list()
                self.draw_annotations()
                self.status_var.set("已删除选中标注")
        else:
            messagebox.showwarning("警告", "请先选择要删除的标注")
            
    def clear_annotations(self):
        """清空所有标注"""
        if self.annotations:
            if messagebox.askyesno("确认", "确定要清空所有标注吗？"):
                self.annotations = []
                self.update_annotation_list()
                self.draw_annotations()
                self.status_var.set("已清空所有标注")
        else:
            messagebox.showinfo("提示", "没有标注需要清空")
            
    def save_annotations(self):
        """保存标注为YOLO格式"""
        if not self.current_image_path or not self.annotations:
            messagebox.showwarning("警告", "没有标注需要保存")
            return
            
        try:
            # 生成标注文件路径
            base_name = os.path.splitext(self.current_image_path)[0]
            annotation_file = base_name + ".txt"
            
            # 获取图像尺寸
            img_width, img_height = self.current_image.size
            
            # 写入YOLO格式标注
            with open(annotation_file, 'w', encoding='utf-8') as f:
                for annotation in self.annotations:
                    class_id = annotation['class_id']
                    annotation_type = annotation.get('type', 'rectangle')
                    
                    if annotation_type == 'rectangle':
                        bbox = annotation['bbox']
                        # 转换为YOLO格式 (归一化坐标)
                        x_center = (bbox[0] + bbox[2]) / 2 / img_width
                        y_center = (bbox[1] + bbox[3]) / 2 / img_height
                        width = (bbox[2] - bbox[0]) / img_width
                        height = (bbox[3] - bbox[1]) / img_height
                        
                        f.write(f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n")
                        
                    elif annotation_type == 'polygon':
                        points = annotation['points']
                        # 多边形格式：class_id x1 y1 x2 y2 ... xn yn
                        line = f"{class_id}"
                        for x, y in points:
                            norm_x = x / img_width
                            norm_y = y / img_height
                            line += f" {norm_x:.6f} {norm_y:.6f}"
                        f.write(line + "\n")
                    
            self.status_var.set(f"已保存标注: {os.path.basename(annotation_file)}")
            messagebox.showinfo("成功", f"标注已保存到: {annotation_file}")
            
        except Exception as e:
            messagebox.showerror("错误", f"保存标注失败: {str(e)}")
            
    def load_annotations(self):
        """加载标注文件"""
        if not self.current_image_path:
            messagebox.showwarning("警告", "请先打开图像")
            return
            
        file_path = filedialog.askopenfilename(
            title="选择标注文件",
            filetypes=[("文本文件", "*.txt"), ("所有文件", "*.*")]
        )
        
        if file_path:
            try:
                self.annotations = []
                with open(file_path, 'r', encoding='utf-8') as f:
                    for line in f:
                        parts = line.strip().split()
                        if len(parts) >= 5:
                            class_id = int(parts[0])
                            x_center = float(parts[1])
                            y_center = float(parts[2])
                            width = float(parts[3])
                            height = float(parts[4])
                            
                            # 转换为像素坐标
                            img_width, img_height = self.current_image.size
                            x1 = (x_center - width/2) * img_width
                            y1 = (y_center - height/2) * img_height
                            x2 = (x_center + width/2) * img_width
                            y2 = (y_center + height/2) * img_height
                            
                            self.annotations.append({
                                'class_id': class_id,
                                'class_name': self.classes[class_id] if class_id < len(self.classes) else f"class_{class_id}",
                                'bbox': (x1, y1, x2, y2)
                            })
                            
                self.update_annotation_list()
                self.draw_annotations()
                self.status_var.set(f"已加载标注: {os.path.basename(file_path)}")
                
            except Exception as e:
                messagebox.showerror("错误", f"加载标注失败: {str(e)}")
                
    def on_key_press(self, event):
        """键盘事件处理"""
        if event.keysym == "Delete":
            self.delete_selected()
        elif event.keysym == "Escape":
            self.canvas.delete("current_bbox")
            self.canvas.delete("polygon_preview")
            self.drawing = False
            if self.annotation_mode == "polygon":
                self.polygon_points = []
                self.status_var.set("已取消多边形绘制")
        elif event.keysym == "s" and event.state & 0x4:  # Ctrl+S
            self.save_annotations()
        elif event.keysym == "o" and event.state & 0x4:  # Ctrl+O
            self.open_image()
        elif event.keysym == "Return" and self.annotation_mode == "polygon":  # Enter完成多边形
            self.finish_polygon()


def main():
    """主函数"""
    root = tk.Tk()
    app = YOLOAnnotator(root)
    root.mainloop()


if __name__ == "__main__":
    main()
